import pandas as pd
from scipy.io import arff
from sklearn.model_selection import train_test_split
from sklearn.preprocessing import StandardScaler
from sklearn.impute import SimpleImputer
from sklearn.ensemble import RandomForestClassifier, StackingClassifier
from sklearn.linear_model import LogisticRegression
from sklearn.metrics import classification_report, confusion_matrix
from imblearn.over_sampling import SMOTE
import matplotlib.pyplot as plt
import seaborn as sns
import numpy as np
from util.LoadData import load_data
from charts.Draw import *

# 预处理数据，包括处理缺失值和转换标签
def preprocess_data(df):
    imputer = SimpleImputer(strategy='mean')
    df_imputed = pd.DataFrame(imputer.fit_transform(df.select_dtypes(include=['number'])), columns=df.select_dtypes(include=['number']).columns)
    df_imputed['Defective'] = df['Defective'].apply(lambda x: 1 if x == 'Y' else 0)
    return df_imputed

# 划分数据集为特征和标签，并进行标准化和数据集划分
def split_data(df):
    X = df.drop('Defective', axis=1)
    y = df['Defective']
    scaler = StandardScaler()
    X_scaled = scaler.fit_transform(X)
    X_train, X_test, y_train, y_test = train_test_split(X_scaled, y, test_size=0.3, random_state=42)
    return X_train, X_test, y_train, y_test

# 应用SMOTE进行过采样
def apply_smote(X_train, y_train, k_value=55):
    smote = SMOTE(k_neighbors=k_value, random_state=42)
    X_train_smote, y_train_smote = smote.fit_resample(X_train, y_train)
    return X_train_smote, y_train_smote

# 训练基础模型
def train_base_models(X_train, y_train):
    rf_model = RandomForestClassifier(n_estimators=100, random_state=42)
    rf_model.fit(X_train, y_train)

    lr_model = LogisticRegression(solver='liblinear', max_iter=10000)
    lr_model.fit(X_train, y_train)

    return rf_model, lr_model

# 训练堆叠模型
def train_stacking_model(file_path):
    df = load_data(file_path)
    df_imputed = preprocess_data(df)
    X_train, X_test, y_train, y_test = split_data(df_imputed)
    X_train, y_train = apply_smote(X_train, y_train)
    rf_model, lr_model = train_base_models(X_train, y_train)
    X_train_rf = rf_model.predict_proba(X_train)
    X_train_lr = lr_model.predict_proba(X_train)

    estimators = [
        ('rf', RandomForestClassifier(n_estimators=100, random_state=42)),
        ('lr', LogisticRegression(solver='liblinear', max_iter=10000))
    ]
    stack_model = StackingClassifier(estimators=estimators, final_estimator=LogisticRegression())
    stack_model.fit(np.hstack([X_train_rf, X_train_lr]), y_train)

    X_test_rf = rf_model.predict_proba(X_test)
    X_test_lr = lr_model.predict_proba(X_test)
    X_test_stack = np.hstack([X_test_rf, X_test_lr])
    return stack_model, X_test_stack, y_test

# 进行预测
def predict_stacking_model(model, X_test_stack):
    y_pred = model.predict(X_test_stack)
    y_pred_prob = model.predict_proba(X_test_stack)[:, 1]
    return y_pred, y_pred_prob

# 展示结果
def show_prediction_results(y_test, y_pred, y_pred_prob):
    report = classification_report(y_test, y_pred, output_dict=True)
    conf_matrix = confusion_matrix(y_test, y_pred)

    metrics_0 = {
        'Accuracy': report['accuracy'],
        'Precision_0': report['0']['precision'],
        'Recall_0': report['0']['recall'],
        'F1_Score_0': report['0']['f1-score']
    }
    plot_metrics(metrics_0, 'Classification Metrics for Class 0 (No Defect)', 'picture/Stacking/classification_metrics_0.png')

    metrics_1 = {
        'Accuracy': report['accuracy'],
        'Precision_1': report['1']['precision'],
        'Recall_1': report['1']['recall'],
        'F1_Score_1': report['1']['f1-score']
    }
    plot_metrics(metrics_1, 'Classification Metrics for Class 1 (Defect)', 'picture/Stacking/classification_metrics_1.png')

    plot_roc_curve(y_test, y_pred_prob, 'picture/Stacking/roc_curve_and_auc.png')
    plot_confusion_matrix(conf_matrix, 'picture/Stacking/confusion_matrix.png')

# 主函数，用于执行整个流程
def main():
    file_path = r'D:\EDAG下载\NASADefectDataset-master\OriginalData\MDP\PC5.arff'
    model, X_test_stack, y_test = train_stacking_model(file_path)
    y_pred, y_pred_prob = predict_stacking_model(model, X_test_stack)
    show_prediction_results(y_test, y_pred, y_pred_prob)

if __name__ == '__main__':
    main()